package com.xialuo.study.std;

import com.google.common.collect.Maps;
import org.apache.commons.lang.StringUtils;
import org.bsc.langgraph4j.serializer.Serializer;
import org.springframework.ai.chat.messages.AssistantMessage;

import java.io.IOException;
import java.io.ObjectInput;
import java.io.ObjectOutput;
import java.util.List;

public class AssistantMessageSerializer implements Serializer<AssistantMessage> {


    @Override
    public void write(AssistantMessage object, ObjectOutput out) throws IOException {
        boolean hasTool = object.hasToolCalls();
        boolean hasText = StringUtils.isNotBlank(object.getText());
        if (!hasText && !hasTool) {
            throw new IllegalArgumentException("Unsupported empty assistant message");
        }

        out.writeBoolean(hasTool);
        out.writeBoolean(hasText);

        if (StringUtils.isNotBlank(object.getText())) {
            out.writeUTF(object.getText());
        }

        if (object.hasToolCalls()) {
            out.writeObject(object.getToolCalls());
        }
    }

    @Override
    public AssistantMessage read(ObjectInput in) throws IOException, ClassNotFoundException {
        boolean hasTool = in.readBoolean();
        boolean hasText = in.readBoolean();
        String text = null;
        if (hasText) {
            text = in.readUTF();
        }

        if (hasTool) {
            List<AssistantMessage.ToolCall> toolCalls = (List<AssistantMessage.ToolCall>)in.readObject();
            return new AssistantMessage(text, Maps.newHashMap(), toolCalls);
        }
        return new AssistantMessage(text);
    }
}
